import os
import json
import re
import time
import traceback
from pydub import AudioSegment
from typing import List
from moviepy.editor import *
from PIL import Image
import natsort
import math
import numpy as np
from glob import glob
import subprocess
from utils import create_directory,read_text_from_file, save_file,load_config
import matplotlib
import matplotlib.pyplot as plt

# plt.ioff()  # 关闭交互模式
# # 您的绘图代码
# plt.show()  # 显式调用 show 来显示图像
# # matplotlib.use('Agg')  # 使用非交互式后端
#
project_config = load_config("config.json")
title = project_config["title"]

def calculate_audio_durations(directory):
    """
    计算指定目录下所有以 audio_for_paragraph_{index} 命名的文件夹中 mp3 文件的总持续时间（以秒为单位）。

    参数:
        directory (str): 需要扫描的根目录路径。

    返回:
        list: 每个 audio_for_paragraph_{index} 文件夹中 mp3 文件总持续时间（秒）的列表。
    """
    # 初始化结果列表
    durations = []

    # 遍历目录下的所有子目录
    for entry in os.scandir(directory):
        if entry.is_dir() and entry.name.startswith("audio_for_paragraph_"):
            # 提取 index
            index = int(entry.name.split("_")[-1])

            # 初始化当前文件夹的总持续时间为0
            total_duration_ms = 0

            # 遍历子目录中的所有文件
            for file_entry in os.scandir(entry.path):
                if file_entry.name.endswith(".mp3"):
                    # 加载 mp3 文件并计算持续时间
                    audio = AudioSegment.from_mp3(file_entry.path)
                    delay = 300
                    total_duration_ms += len(audio) + delay

            # 将当前文件夹的总持续
            # 时间转换为秒，并添加到结果列表中
            total_duration_seconds = total_duration_ms / 1000.0
            durations.append((index, total_duration_seconds))

    # 按照 index 排序结果列表
    durations.sort(key=lambda x: x[0])



    # 只保留持续时间（秒）
    durations = [duration for _, duration in durations]

    durations.insert(0, 2)

    return durations

# 计算各段落的所有音频时长
audio_file_folder = project_config["audio_file_folder"].format(title=project_config["title"])

# 计算音频时长
durations = calculate_audio_durations(audio_file_folder)

durations_file=project_config["durations_folder"].format(title=project_config["title"])

create_directory(durations_file)

# 打印结果
print("各段落的音频时长（秒）:")
with open(durations_file, "w") as f:
    for index, duration in enumerate(durations):
        f.write(f"段落 {index + 1}: {duration:.2f} 秒\n")
        print(f"段落 {index + 1}: {duration:.2f} 秒")

print(f"时长信息已保存到 {durations_file}")

def images_to_video_with_durations(input_image_path, output_video_path, durations, fps, base_name):
    # 获取所有符合条件的图片，并按文件名中的数字排序
    # pattern = r'^' + re.escape(base_name) + r'_(\d+)\.png$'
    # pattern = r".*_(\d+)\.png"
    pattern = r".*_((?:\d+))_.*\.png$"
    image_files = [
        f"{input_image_path}/{file}"
        for file in os.listdir(input_image_path)
        if re.match(pattern, file)
    ]

    print("Matching files:", image_files)  # 调试输出，查看匹配的文件
    print("fps=", fps)

    image_files = natsort.natsorted(image_files, key=lambda x: int(re.match(pattern, os.path.basename(x)).group(1)))

    # 确定视频的背景尺寸
    target_width, target_height = 1280, 720
    background_size = (target_width, target_height)

    clips = []
    for i, file in enumerate(image_files):
        print(f"Processing file: {file}, duration: {durations[i]}")  # 再次调试输出
        img = Image.open(file)
        width, height = img.size
        ratio = width / height

        if width > target_width or height > target_height:
            if ratio > target_width / target_height:
                new_width = target_width
                new_height = math.floor(new_width / ratio)
            else:
                new_height = target_height
                new_width = math.floor(new_height * ratio)
        else:
            new_width, new_height = width, height

        img = img.resize((new_width, new_height), resample=Image.Resampling.LANCZOS)
        img_clip = ImageClip(np.array(img)).set_duration(durations[i])
        img_clip = img_clip.set_position('center')
        bg_clip = ColorClip(size=background_size, color=(255, 255, 255), duration=durations[i])

        composite_clip = CompositeVideoClip([bg_clip, img_clip])
        composite_clip = composite_clip.set_fps(30)
        clips.append(composite_clip)

    # 使用concatenate_videoclips函数将所有剪辑串联在一起
    final_clip = concatenate_videoclips(clips, method="compose")
    output_filename = os.path.join(output_video_path, f"{base_name}.mp4")
    print('outputfilename=', output_filename)
    create_directory(output_filename)
    print('fps=', fps)
    # 确保 final_clip 有一个有效的 fps 属性
    if final_clip.fps is None:
        final_clip = final_clip.set_fps(30)
    print('fps=', final_clip.fps)

    # todo 源代码需要修改fps =30
    final_clip.write_videofile(output_filename, fps=30)

def merge_audio_and_add_to_video(video_path, audio_base_dir, output_path):
    """
    合并多个音频文件并添加到视频中。

    :param video_path: 视频文件的路径。
    :param audio_base_dir: 包含音频文件夹的基目录。
    :param output_path: 输出视频的路径。
    """
    # 加载视频文件
    video_clip = VideoFileClip(video_path)

    # 初始化音频列表
    audio_clips = []

    silent_audio_start = AudioClip(lambda t: [0,0], duration=2)
    audio_clips.append(silent_audio_start)

    # 遍历所有子目录，按数字大小排序
    audio_dirs = glob(os.path.join(audio_base_dir, "audio_for_paragraph_*"))
    audio_dirs.sort(key=lambda x: int(re.search(r'\d+', os.path.basename(x)).group()))

    # 遍历所有子目录
    for audio_dir in audio_dirs:
        # 获取当前目录的index
        index = int(os.path.basename(audio_dir).split("_")[-1])

        # 遍历目录中的所有mp3文件
        mp3_files = glob(os.path.join(audio_dir, f"paragraph_{index}_sentence_*.mp3"))
        mp3_files.sort(key=lambda x: int(re.search(r'_sentence_(\d+)', os.path.basename(x)).group(1)))

        # 遍历排序后的mp3文件列表
        for mp3_file in mp3_files:
            # 加载音频文件
            audio_clip = AudioFileClip(mp3_file)

            # 添加到音频列表
            if audio_clips:
                # 如果不是第一个音频，则在前一个音频之后添加0.5秒的静音
                # 替换原有的 AudioNullClip 代码
                silent_audio = AudioClip(lambda t: [0,0], duration=0.3)
                audio_clips.append(silent_audio)
            audio_clips.append(audio_clip)

    # 合并所有音频片段
    final_audio = concatenate_audioclips(audio_clips)

    # 将音频添加到视频中
    video_with_audio = video_clip.set_audio(final_audio)

    # 输出带有新音频的视频文件
    video_with_audio.write_videofile(output_path, fps=30, codec='libx264', audio_codec='aac')

    # 关闭剪辑对象，释放资源
    video_clip.close()

def merge_video_and_subtitle(video_path, srt_path, output_path):
    # 如果输出文件已存在，删除它
    if os.path.exists(output_path):
        os.remove(output_path)

    command = [
        'ffmpeg',
        '-i', video_path,
        '-vf', f'subtitles={srt_path}',
        '-c:a', 'copy',
        output_path
    ]

    subprocess.run(command, check=True)


marp_export_image_folder = project_config["marp_export_image_folder"].format(title=project_config["title"])
srt_and_video_folder = project_config["srt_and_video_folder"]
fps =  project_config["fps"]

marp_export_image_folder = './output/illustrations'

# 检查输入图像
# pattern = r'^' + re.escape(input_base_name) + r'_(\d+)\.png$'
# pattern = r".*_(\d+)\.png"
pattern = r".*_((?:\d+))_.*\.png$"

# 获取目录中的所有文件
all_files = os.listdir(marp_export_image_folder)

# 过滤出符合模式的PNG文件，并构建完整路径
# image_files = [
#     os.path.join(marp_export_image_folder, file)
#     for file in all_files
#     if re.match(pattern, file) and file.lower().endswith('.png')
# ]


image_files = [
    f"{marp_export_image_folder}/{file}"
    for file in os.listdir(marp_export_image_folder)
    if re.match(pattern, file)
]

# 打印调试信息
print(f"所有文件: {all_files}")
print(f"匹配到的图像文件: {image_files}")

print("Matching files:", image_files)  # 输出匹配的文件

if not image_files:
    raise ValueError("No matching image files found.")

# 检查 durations 的数量
if len(durations) != len(image_files):
    raise ValueError("The number of durations must match the number of image files.")

# 调用函数
images_to_video_with_durations(
    marp_export_image_folder,
    srt_and_video_folder,
    durations,
    fps,
    project_config["title"]
)


# 合成路径
video_raw = project_config["video_raw"].format(title=project_config["title"])# 视频文件的路径
audio_file_folder = project_config["audio_file_folder"].format(title=project_config["title"])
# audio_base_dir = "./output/audio/"+title+"_课程脚本_speech_script_plus"  # 音频文件夹的基目录
video_with_audio = project_config["video_with_audio"].format(title=project_config["title"]) # 输出视频的路径

# 检查视频和音频路径是否存在
if not os.path.exists(video_raw):
    raise ValueError(f"Video file not found: {video_raw}")

if not os.path.exists(audio_file_folder):
    raise ValueError(f"Audio directory not found: {audio_file_folder}")

# 调用函数
merge_audio_and_add_to_video(video_raw, audio_file_folder, video_with_audio)


# 构建文件路径
video_with_audio = project_config["video_with_audio"].format(title=project_config["title"])
srt_file_path = project_config["srt_file_path"].format(title=project_config["title"])

video_with_audio_and_subtitles = project_config["video_with_audio_and_subtitles"].format(title=project_config["title"])

# 调用函数
merge_video_and_subtitle(video_with_audio, srt_file_path, video_with_audio_and_subtitles)

from IPython.display import Video

# 本地视频文件路径
video_path = project_config["video_with_audio_and_subtitles"].format(title=project_config["title"])
print("视频地址：",video_path)

# 播放视频
Video(video_path, width=768, height=512)